import datetime
import pandas as pd
import tushare as ts
import matplotlib.pyplot as plt
from matplotlib.widgets import MultiCursor
from pandas.plotting import register_matplotlib_converters

import QUANTAXIS as QA


register_matplotlib_converters()


class NorthData():
    def __init__(self, codes, start, end):
        self.codes = codes
        data_stock = QA.QA_fetch_stock_day_adv(codes['code'].to_list(), start, end)
        data_stock = data_stock.to_qfq()
        data_north = QA.QA_fetch_north_stock(code=codes['code'].to_list(), start=start, end=end)
        data_stock.data = pd.merge(data_stock.data, data_north, left_index=True, right_index=True, how='left')
        # stock = pd.merge(data_stock, data_north, on=['date', 'code'], how='left')
        data_stock.data.dropna(axis=0, inplace=True)
        # self.data_stock = stock.set_index(['date', 'code'])
        self.data_stock = data_stock.data
        self.struct_stock = data_stock

        data_index = QA.QA_fetch_index_day_adv(['000001', '000300'], start, end)
        data_north_day = QA.QA_fetch_north_day(start=start, end=end)
        data_index.data = pd.merge(data_index.data, data_north_day, left_index=True, right_index=True, how='left')
        data_index.data.fillna(0, inplace=True)
        self.data_day = data_index.data
        self.struct_day = data_index

    # grow up TOP 5 fastest stocks in N bars
    def ratio_up_series_fast(self, N=5):
        N = N + 1  # the first bar is baseline


    def data_to_csv(self):
        self.data_stock.to_csv('stock.csv', encoding='utf_8_sig')
        self.data_day.to_csv('day.csv', encoding='utf_8_sig')

    def plot_indicator(self, indicator):
        ind = self.struct_stock.add_func(indicator, N=10, M=3)
        for _, code in self.codes.iterrows():
            if code['code'] not in self.struct_stock.index.levels[1]:
                continue
            df = ind.xs(code['code'], level='code', axis=0, drop_level=True)
            ind_last_last_bar = df.iloc[-3]
            ind_last_bar = df.iloc[-2]
            ind_current_bar = df.iloc[-1]
            if ind_current_bar.XC > 0:
                print('code：{}---{}---时间{}---CLOSE{}'.format(
                    code['code'], code['name'], ind_current_bar.name, ind_current_bar.CLOSE))

                figure = plt.figure()
                # plt.title('{}---'.format(code['code']))
                axes1 = figure.add_subplot(411)
                axes2 = figure.add_subplot(412)
                axes3 = figure.add_subplot(413)
                axes4 = figure.add_subplot(414)
                axes1.plot(df.index, df['CLOSE'])
                axes2.plot(df.index, df['VOLUME'])
                axes3.plot(df.index, df[['RATIO']])
                axes4.plot(df.index, df[['XC']])

                # left, right = axes2.get_xlim()
                # axes2.hlines(y=0, xmin=left, xmax=right, linestyles='dashed')
                multi = MultiCursor(figure.canvas, (axes1, axes2, axes3, axes4), color='r', lw=1)

                plt.show()

    def plot_north_day(self):
        df = self.data_day.xs('000300', level='code', axis=0, drop_level=False)
        figure = plt.figure()
        axes1 = figure.add_subplot(411)
        axes2 = figure.add_subplot(412)
        axes3 = figure.add_subplot(413)
        axes4 = figure.add_subplot(414)
        axes1.plot(self.data_day.index.levels[0], df['close'])
        axes2.plot(self.data_day.index.levels[0], df['volume'])
        axes3.plot(self.data_day.index.levels[0], df[['ggt_ss', 'ggt_sz']])
        # axes4.plot(self.data_day.index.levels[0], self.data_day[['north_money', 'south_money']])
        axes4.plot(self.data_day.index.levels[0], df[['hgt', 'sgt']])

        left, right = axes3.get_xlim()
        axes3.hlines(y=0, xmin=left, xmax=right, linestyles='dashed')
        axes4.hlines(y=0, xmin=left, xmax=right, linestyles='dashed')
        multi = MultiCursor(figure.canvas, (axes1, axes2, axes3, axes4), color='r', lw=1)
        plt.show()

    def plot_north_stock(self, code='000001'):
        data = self.data_stock.xs(code, level='code', axis=0, drop_level=True)
        figure = plt.figure()
        # plt.title('{}---'.format('code'))
        axes1 = figure.add_subplot(411)
        axes2 = figure.add_subplot(412)
        axes3 = figure.add_subplot(413)
        axes4 = figure.add_subplot(414)
        axes1.plot(data.index, data['close'])
        axes2.plot(data.index, data[['ratio']])
        axes3.bar(data.index, data['ratio'].diff(1))
        axes4.plot(data.index, data['volume'])

        left, right = axes3.get_xlim()
        axes3.hlines(y=0, xmin=left, xmax=right, linestyles='dashed')
        multi = MultiCursor(figure.canvas, (axes1, axes2, axes3, axes4), color='r', lw=1)
        plt.show()


if __name__ == '__main__':
    code = '600085'  # 300604 000738 600893
    stock_list = QA.QA_fetch_stock_list_adv()
    stock_list = stock_list[stock_list.code.str.startswith(code)]
    codes = stock_list[~stock_list.name.str.contains('ST')]
    # codes = pd.read_csv('D:\\PythonPro\\QUANTAXIS\\EXAMPLE\\AI\\data\\hs300.csv', dtype=str)
    start = '2020-12-01'
    # end = str(datetime.date.today() - datetime.timedelta(days=1))
    end = str(datetime.date.today())
    north = NorthData(codes, start, end)
    # north.data_to_csv()
    north.plot_north_stock(code=code)  # 600036 601318
    # north.plot_north_day()
    # north.plot_indicator(QA.QA_indicator_NORTH_XC)
    # for _, code in codes.iterrows():
    #     if code['code'] not in north.struct_stock.index.levels[1]:
    #         continue
    #     print(code['name'])
    #     north.plot_north_stock(code=code['code'])
